#include <stdio.h>
#include <stdlib.h>

#include "include/pf_tree.h"
#include "include/pf_internal.h"

/*! Color enumeration for nodes of red-black tree */
typedef enum _red_black_color_enum
{
    rbcRed,
    rbcBlack
} red_black_color_enum;

struct pf_tree_node
{
    void*                       object;       /* the stored object user defined */
    red_black_color_enum        color;        /* the color of the node */
    struct pf_tree_node*        parent;       /* points to the parent node */
    struct pf_tree_node*        right;        /* points to the right child */
    struct pf_tree_node*        left;         /* points to the left child */
};

struct pf_tree
{
    struct pf_tree_node*        root;         /* pointer to the tree root */
    uint32_t                    iSize;        /* number of objects stored */
    int (*comp)(const void*, const void*);    /* compare function */
};

#define RB_PARENT(node)             ((pf_tree_node_t*)((node)->parent))
#define RB_COLOR(node)              ((node)->color)
#define RB_SET_RED(node)            ((node)->color = rbcRed)
#define RB_SET_BLACK(node)          ((node)->color = rbcBlack)
#define RB_IS_RED(node)             ((node)->color == rbcRed)
#define RB_IS_BLACK(node)           ((node)->color == rbcBlack)
#define RB_SET_COLOR(node, color)   ((node)->color = (color))
#define RB_SET_PARENT(node1, node2) ((node1)->parent = (node2))

/********************************************************************************************/
 
#ifdef UNIT_TEST
int test_get_color(pf_tree_node_t* node)
{
    return (int)RB_COLOR(node);
}

pf_tree_node_t* test_get_parent(pf_tree_node_t* node)
{
    return RB_PARENT(node);
}

pf_tree_node_t* test_get_left(pf_tree_node_t* node)
{
    return node->left;
}

pf_tree_node_t* test_get_right(pf_tree_node_t* node)
{
    return node->right;
}
#endif

/********************************************************************************************/

/*********************************************************************************************
Function Name:  rbnode_construct
Description  :  construct a new tree node
Inputs       :  void* object                :  the tree node's data
                red_black_color_enum color  :  the tree node's color
Outputs      :  return value                :  the tree node's pointer
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline pf_tree_node_t* rbnode_construct(void* object, red_black_color_enum color);

/*********************************************************************************************
Function Name:  rbnode_destruct
Description  :  destruct a tree node
Inputs       :  pf_tree_node_t* node   :  one tree node's pointer
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static void rbnode_destruct(pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbnode_minimum
Description  :  Get the leftmost node in the sub-tree spanned by the given node
Inputs       :  pf_tree_node_t* node        :  The sub-tree root node
Outputs      :  return value                :  The sub-tree minimum node
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline pf_tree_node_t* rbnode_minimum(pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbnode_maximum
Description  :  Get the rightmost node in the sub-tree spanned by the given node
Inputs       :  pf_tree_node_t* node        :  The sub-tree root node
Outputs      :  return value                :  The sub-tree maximum node
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline pf_tree_node_t* rbnode_maximum(pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbtree_rotate_left
Description  :  Left-rotate the sub-tree spanned by the given node
Inputs       :  pf_tree_t* tree             :  The rb tree
                pf_tree_node_t* node        :  The sub-tree node which need to rotate_left
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline void rbtree_rotate_left(pf_tree_t* tree, pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbtree_rotate_right
Description  :  Right-rotate the sub-tree spanned by the given node
Inputs       :  pf_tree_t* tree             :  The rb tree
                pf_tree_node_t* node        :  The sub-tree node which need to rotate_right
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline void rbtree_rotate_right(pf_tree_t* tree, pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbtree_insert_fixup
Description  :  Fix-up the red-black tree properties after an insertion operation
Inputs       :  pf_tree_t* tree             :  The rb tree
                pf_tree_node_t* node        :  The node that has just been inserted to the tree
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline void rbtree_insert_fixup(pf_tree_t* tree, pf_tree_node_t* node);

/*********************************************************************************************
Function Name:  rbtree_remove_fixup
Description  :  Fix-up the red-black tree properties after a removal operation
Inputs       :  pf_tree_t* tree             :  The rb tree
                pf_tree_node_t* node        :  The child of the node that has just been removed from the tree
                pf_tree_node_t* parent      :  The parent node
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline void rbtree_remove_fixup(pf_tree_t *tree, 
                                       pf_tree_node_t *node, 
                                       pf_tree_node_t* parent);

/*********************************************************************************************
Function Name:  rbnode_traverse
Description  :  Traverse a red-black tree left first
Inputs       :  pf_tree_node_t* node                    :  The rb tree
                int  (*cb)(void* arg, void* value)      :  The operation to perform on 
                                                           every object of the tree
                void* arg                               :  The operation 's params
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline int rbnode_traverse(pf_tree_node_t* node, 
                                  int (*cb)(void* arg, pf_tree_node_t* node), 
                                  void* arg);

/*********************************************************************************************
Function Name:  rbnode_traverse_right
Description  :  Traverse a red-black tree right first
Inputs       :  pf_tree_node_t* node                    :  The rb tree
                int  (*cb)(void* arg, void* value)      :  The operation to perform on 
                                                           every object of the tree
                void* arg                               :  The operation 's params
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline int rbnode_traverse_right(pf_tree_node_t* node, 
                                        int (*cb)(void* arg, pf_tree_node_t* node), 
                                        void* arg);

/*********************************************************************************************
Function Name:  tree_erase_impl
Description  :  
Inputs       :  pf_tree_t* tree              :  
                pf_tree_node_t* node         :  
Outputs      :
ErrorCodes   :
History      :
---------------------------------------------------------------------------------------------
Author                 Date                          Comments
zhanggp                2011-06-29                    create
**********************************************************************************************/
static inline void* tree_erase_impl(pf_tree_t* tree, pf_tree_node_t* node);

/********************************************************************************************/
 
pf_tree_t* pf_tree_create(int (*comp)(const void*, const void*))
{
    pf_tree_t* tree = (pf_tree_t*) malloc(sizeof(pf_tree_t));

    if(!tree)
    {
        pf_set_errno(PF_ENOMEM);
        return NULL;
    }
    tree->comp = comp;
    tree->iSize = 0;
    tree->root = NULL;

    return tree;
}

int pf_tree_destory(pf_tree_t* tree)
{
    if(tree->root)
    {
        rbnode_destruct(tree->root);
    }
    free(tree);

    return 0;
}

void* pf_tree_get_value(const pf_tree_node_t* node)
{
    assert(node);

    return node->object;
}

pf_tree_node_t* pf_tree_head(const pf_tree_t* tree)
{
    assert(tree);

    if(tree->root == NULL)
    {
        return NULL;
    }
    return rbnode_minimum(tree->root);
}

pf_tree_node_t* pf_tree_tail(const pf_tree_t* tree)
{
    assert(tree);

    if(tree->root == NULL)
    {
        return NULL;
    }
    return rbnode_maximum(tree->root);
}

pf_tree_node_t* pf_tree_next(const pf_tree_node_t* cur_node)
{
    pf_tree_node_t* succ_node       = NULL;
    const pf_tree_node_t* prev_node = NULL;

    assert(cur_node);

    if(cur_node->right)
    {
        /* If there is a right child, the successor is the minimal object in
           the sub-tree spanned by this child.  */
        succ_node = cur_node->right;
        while(succ_node->left)
        {
            succ_node = succ_node->left;
        }
    }
    else
    {
        /* Otherwise, go up the tree until reaching the parent from the left
           direction.  */
        prev_node = cur_node;
        succ_node = cur_node->parent;
        while(succ_node && prev_node == succ_node->right)
        {
            prev_node = succ_node;
            succ_node = succ_node->parent;
        }
    }

    return succ_node;
}

pf_tree_node_t* pf_tree_prev(const pf_tree_node_t* cur_node)
{
    pf_tree_node_t* pred_node       = NULL;
    const pf_tree_node_t* prev_node = NULL;

    assert(cur_node);

    if(cur_node->left)
    {
        /* If there is a left child, the predecessor is the maximal object in
           the sub-tree spanned by this child.  */
        pred_node = cur_node->left;
        while(pred_node->right)
        {
            pred_node = pred_node->right;
        }
    }
    else
    {
        /* Otherwise, go up the tree until reaching the parent from the right
           direction. */
        prev_node = cur_node;
        pred_node = cur_node->parent;
        while(pred_node && prev_node == pred_node->left)
        {
            prev_node = pred_node;
            pred_node = pred_node->parent;
        }
    }

    return pred_node;
}

pf_tree_node_t* pf_tree_find(const pf_tree_t* tree, const void* value)
{
    pf_tree_node_t* cur_node = tree->root;
    int comp_result               = 0;

    assert(tree);
    assert(value);

    while(cur_node)
    {
        /* In case of equality, we can return the current node. */
        if((comp_result = tree->comp(cur_node->object, value)) == 0)
        {
            return cur_node;
        }
        /* Go down to the left or right child. */
        cur_node = (comp_result > 0) ? cur_node->left : cur_node->right;
    }
    /* If we reached here, the object is not found in the tree */
    return NULL;
}

pf_tree_node_t* pf_tree_find_less(const pf_tree_t* tree, 
                                       const void* value, 
                                       int can_equal)
{
    pf_tree_node_t* cur_node    = tree->root;
    pf_tree_node_t* last_node   = NULL;
    int comp_result                  = 0;

    assert(tree);
    assert(value);

    while(cur_node)
    {
        comp_result = tree->comp(cur_node->object, value);
        if(comp_result >= 0)
        {
            if((can_equal != 0) && (comp_result == 0))
            {
                return cur_node;
            }
            cur_node = cur_node->left;
        }
        else
        {
            last_node = cur_node;
            cur_node = cur_node->right;
        }
    }

    return last_node;
}


int pf_tree_insert(pf_tree_t* tree, void* value)
{
    int cmp                       = 0;
    pf_tree_node_t* cur_node = NULL;
    pf_tree_node_t* new_node = NULL;

    assert(tree);
    assert(value);

    if(tree->iSize == UINT32_MAX)
    {
        pf_set_errno(PF_EDQUOT);
        return PF_EDQUOT;    
    }

    if(!(tree->root))
    {
        /* Assign a new root node. Notice that the root is always black */
        new_node = rbnode_construct(value, rbcBlack);
        if(!new_node)
        {
            pf_set_errno(PF_ENOMEM);
            return PF_ENOMEM;
        }
        tree->root = new_node;
        tree->iSize = 1;

        return 0;
    }
    /* Find a place for the new object, and insert it as a red leaf */
    cur_node = tree->root;

    while(cur_node)
    {
        cmp = tree->comp(cur_node->object, value);
        /* Compare inserted object with the object stored in the current node */
        if(cmp > 0)
        {
            if(!(cur_node->left))
            {
                /* Insert the new leaf as the left child of the current node */
                new_node = rbnode_construct(value, rbcRed);
                if(!new_node)
                {
                    pf_set_errno(PF_ENOMEM);
                    return PF_ENOMEM;
                }
                cur_node->left = new_node;
                new_node->parent = cur_node;
                cur_node = NULL;                /* terminate the while loop */
            }
            else
            {
                cur_node = cur_node->left;      /* Go to the left sub-tree */
            }
        }
        else if(cmp < 0)
        {
            if(!(cur_node->right))
            {
                /* Insert the new leaf as the right child of the current node */
                new_node = rbnode_construct(value, rbcRed);
                if(!new_node)
                {
                    pf_set_errno(PF_ENOMEM);
                    return PF_ENOMEM;
                }
                cur_node->right = new_node;
                new_node->parent = cur_node;
                cur_node = NULL;                /* terminate the while loop */
            }
            else
            {
                cur_node = cur_node->right;     /* Go to the right sub-tree */
            }
        }
        else
        {
            pf_set_errno(PF_EKEYREJECTED);
            return PF_EKEYREJECTED;
        }
    }
    /* Mark that a new node was added */
    tree->iSize++;
    /* Fix up the tree properties */
    rbtree_insert_fixup(tree, new_node);

    return 0;
}


void* pf_tree_erase(pf_tree_t* tree, pf_tree_node_t* node)
{
    assert(tree);
    assert(node);

    return tree_erase_impl(tree, node);
}

int pf_tree_erase_value(pf_tree_t* tree, const void* value)
{
    pf_tree_node_t* node = NULL;
    assert(tree);
    assert(value);

    node = pf_tree_find(tree, value);
    if(!node)
    {
        pf_set_errno(PF_ENODATA);
        return PF_ENODATA;
    }
    tree_erase_impl(tree, node);

    return 0;
}

uint32_t pf_tree_count(const pf_tree_t* tree)
{
    assert(tree);

    return tree->iSize;
}

int pf_tree_clear(pf_tree_t* tree)
{
    assert(tree);

    if(tree->root)
    {
        rbnode_destruct(tree->root);
    }

    tree->root = NULL;
    tree->iSize = 0;

    return 0;
}

int32_t pf_tree_for_each(const pf_tree_t* tree,
                         int forward,
                         void* arg,
                         int (*cb)(void* arg, pf_tree_node_t* node))
{
    assert(tree);
    assert(cb);

    if(forward)
    {
        return rbnode_traverse(tree->root, cb,arg);
    }
    else
    {
        return rbnode_traverse_right(tree->root, cb,arg);
    }

    return 0;
}

/********************************************************************************************/

static inline pf_tree_node_t* rbnode_construct(void* object, red_black_color_enum color)
{
    pf_tree_node_t* node = (pf_tree_node_t*) malloc(sizeof(pf_tree_node_t));

    if(!node)
    {
        return NULL;
    }
    node->object = object;
    node->color = color;
    node->parent = node->right = node->left = NULL;

    return node;
}

static void rbnode_destruct(pf_tree_node_t* node)
{
    if(!node) 
    {
        return;
    }

    rbnode_destruct(node->right);
    rbnode_destruct(node->left);
    free(node);

    return;
}

static inline pf_tree_node_t* rbnode_minimum(pf_tree_node_t* node)
{
    while(node->left)
    {
        node = node->left;
    }

    return node;
}

static inline pf_tree_node_t* rbnode_maximum(pf_tree_node_t* node)
{
    while(node->right)
    {
        node = node->right;
    }

    return node;
}

static inline void rbtree_rotate_left(pf_tree_t* tree, pf_tree_node_t* node)
{
    pf_tree_node_t* right = node->right;
    pf_tree_node_t* parent = RB_PARENT(node);

    if((node->right = right->left))
    {
        RB_SET_PARENT(right->left, node);
    }
    right->left = node;

    RB_SET_PARENT(right, parent);

    if(parent)
    {
        if(node == parent->left)
        {
            parent->left = right;
        }
        else
        {
            parent->right = right;
        }
    }
    else
    {
        tree->root = right;
    }
    RB_SET_PARENT(node, right);

    return;
}

static inline void rbtree_rotate_right(pf_tree_t* tree, pf_tree_node_t* node)
{
    pf_tree_node_t *left = node->left;
    pf_tree_node_t *parent = RB_PARENT(node);

    if((node->left = left->right))
    {
        RB_SET_PARENT(left->right, node);
    }
    left->right = node;

    RB_SET_PARENT(left, parent);

    if(parent)
    {
        if(node == parent->right)
        {
            parent->right = left;
        }
        else
        {
            parent->left = left;
        }
    }
    else
    {
        tree->root = left;
    }
    RB_SET_PARENT(node, left);

    return;
}

static inline void rbtree_insert_fixup(pf_tree_t* tree, pf_tree_node_t* node)
{
    pf_tree_node_t* parent = NULL;
    pf_tree_node_t* gparent = NULL;

    while((parent = RB_PARENT(node)) && RB_IS_RED(parent))
    {
        gparent = RB_PARENT(parent);

        if(parent == gparent->left)
        {
            {
                register pf_tree_node_t* uncle = gparent->right;
                if(uncle && RB_IS_RED(uncle))
                {
                    RB_SET_BLACK(uncle);
                    RB_SET_BLACK(parent);
                    RB_SET_RED(gparent);
                    node = gparent;
                    continue;
                }
            }

            if(parent->right == node)
            {
                register pf_tree_node_t* tmp;
                rbtree_rotate_left(tree, parent);
                tmp = parent;
                parent = node;
                node = tmp;
            }

            RB_SET_BLACK(parent);
            RB_SET_RED(gparent);
            rbtree_rotate_right(tree, gparent);
        } else {
            {
                register pf_tree_node_t* uncle = gparent->left;
                if(uncle && RB_IS_RED(uncle))
                {
                    RB_SET_BLACK(uncle);
                    RB_SET_BLACK(parent);
                    RB_SET_RED(gparent);
                    node = gparent;
                    continue;
                }
            }

            if(parent->left == node)
            {
                register pf_tree_node_t* tmp;
                rbtree_rotate_right(tree, parent);
                tmp = parent;
                parent = node;
                node = tmp;
            }

            RB_SET_BLACK(parent);
            RB_SET_RED(gparent);
            rbtree_rotate_left(tree, gparent);
        }
    }

    RB_SET_BLACK(tree->root);

    return;
}

static inline void rbtree_remove_fixup(pf_tree_t *tree, pf_tree_node_t *node, pf_tree_node_t* parent)
{
    pf_tree_node_t* other = NULL;
    red_black_color_enum color;

    while((!node || RB_IS_BLACK(node)) && node != tree->root)
    {
        if(parent->left == node)
        {
            other = parent->right;
            if(RB_IS_RED(other))
            {
                RB_SET_BLACK(other);
                RB_SET_RED(parent);
                rbtree_rotate_left(tree, parent);
                other = parent->right;
            }
            if((!other->left || RB_IS_BLACK(other->left)) 
               && (!other->right || RB_IS_BLACK(other->right)))
            {
                RB_SET_RED(other);
                node = parent;
                parent = RB_PARENT(node);
            }
            else
            {
                if(!other->right || RB_IS_BLACK(other->right))
                {
                    RB_SET_BLACK(other->left);
                    RB_SET_RED(other);
                    rbtree_rotate_right(tree, other);
                    other = parent->right;
                }
                color = RB_COLOR(parent);
                RB_SET_COLOR(other, color);
                RB_SET_BLACK(parent);
                RB_SET_BLACK(other->right);
                rbtree_rotate_left(tree, parent);
                node = tree->root;
                break;
            }
        }
        else
        {
            other = parent->left;
            if(RB_IS_RED(other))
            {
                RB_SET_BLACK(other);
                RB_SET_RED(parent);
                rbtree_rotate_right(tree, parent);
                other = parent->left;
            }
            if((!other->left || RB_IS_BLACK(other->left)) 
               && (!other->right || RB_IS_BLACK(other->right)))
            {
                RB_SET_RED(other);
                node = parent;
                parent = RB_PARENT(node);
            }
            else
            {
                if(!other->left || RB_IS_BLACK(other->left))
                {
                    RB_SET_BLACK(other->right);
                    RB_SET_RED(other);
                    rbtree_rotate_left(tree, other);
                    other = parent->left;
                }
                color = RB_COLOR(parent);
                RB_SET_COLOR(other, color);
                RB_SET_BLACK(parent);
                RB_SET_BLACK(other->left);
                rbtree_rotate_right(tree, parent);
                node = tree->root;
                break;
            }
        }
    }
    if(node)
    {
        RB_SET_BLACK(node);
    }

    return;
}

static inline int rbnode_traverse(pf_tree_node_t* node, 
                                  int (*cb)(void* arg, pf_tree_node_t* node), 
                                  void* arg)
{
    int ret = 0;

    if(!node)
    {
        return 0;
    }
    if((ret = rbnode_traverse(node->left, cb, arg)) != 0)
    {
        return ret;
    }
    pf_tree_node_t* cur = node;
    node = cur->right;
    if((ret = cb(arg, cur)) != 0)
    {
        return ret;
    }
    if((ret = rbnode_traverse(node, cb, arg)) != 0)
    {
        return ret;
    }

    return 0;
}

static inline int rbnode_traverse_right(pf_tree_node_t* node, 
                                        int (*cb)(void* arg, pf_tree_node_t* node), 
                                        void* arg)
{
    int ret = 0;

    if(!node)
    {
        return 0;
    }
    if((ret = rbnode_traverse_right(node->right, cb, arg)) != 0)
    {
        return ret;
    }
    pf_tree_node_t* cur = node;
    node = cur->left;
    if((ret = cb(arg, cur)) != 0)
    {
        return ret;
    }
    if((ret = rbnode_traverse_right(node, cb, arg)) != 0)
    {
       return ret;
    }

    return 0;
}

static inline void* tree_erase_impl(pf_tree_t* tree,pf_tree_node_t* node)
{
    void *value = NULL;
    pf_tree_node_t* cur_node = node;
    pf_tree_node_t* child = NULL;
    pf_tree_node_t* parent = NULL;
    red_black_color_enum color;

    if(!node->left)
    {
        child = node->right;
    }
    else if(!node->right)
    {
        child = node->left;
    }
    else
    {
        pf_tree_node_t *old = node, *left;

        node = node->right;
        while((left = node->left) != NULL)
            node = left;

        if(RB_PARENT(old)) 
        {
            if(RB_PARENT(old)->left == old)
            {
                RB_PARENT(old)->left = node;
            }
            else
            {
                RB_PARENT(old)->right = node;
            }
        } 
        else
        {
            tree->root = node;
        }

        child = node->right;
        parent = RB_PARENT(node);
        color = RB_COLOR(node);

        if(parent == old) 
        {
            parent = node;
        } 
        else 
        {
            if(child)
            {
                RB_SET_PARENT(child, parent);
            }
            parent->left = child;

            node->right = old->right;
            RB_SET_PARENT(old->right, node);
        }

        node->parent = old->parent;
        node->color = old->color;
        node->left = old->left;
        RB_SET_PARENT(old->left, node);

        goto color;
    }

    parent = RB_PARENT(node);
    color = RB_COLOR(node);

    if(child)
    {
        RB_SET_PARENT(child, parent);
    }
    if(parent)
    {
        if(parent->left == node)
        {
            parent->left = child;
        }
        else
        {
            parent->right = child;
        }
    }
    else
    {
        tree->root = child;
    }

color:
    if(color == rbcBlack)
    {
        rbtree_remove_fixup(tree, child, parent);
    }
    cur_node->left = NULL;
    cur_node->right = NULL;
    value = cur_node->object;
    free(cur_node);
    
    tree->iSize--;
    return value;
}
